# *-------------检测视频，将视频切帧后进行逐帧检测，检测完后再合并成新视频------------------*
import argparse
import time
from pathlib import Path

import cv2
from PIL import Image
import numpy as np
import os
import torch
import torch.backends.cudnn as cudnn
from numpy import random
from tqdm import tqdm

from models.experimental import attempt_load
from utils.datasets import LoadStreams, LoadImages
from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
    scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path
from utils.plots import plot_one_box, amount_name
from utils.torch_utils import select_device, load_classifier, time_synchronized

name = []  # 记录类别名字
color = []  # 记录这个类别名对应的颜色

first_path = 'D://PyCharm/Code/yolov5-5.0/yolov5-5.0/'
video_path = 'D://PyCharm/Code/yolov5-5.0/yolov5-5.0/data/video/8.mp4'
frame_occasional_save = 'D://PyCharm/Code/yolov5-5.0/yolov5-5.0/data/images/frames/'

video_result_path = 'D://PyCharm/Code/yolov5-5.0/yolov5-5.0/data/video/result/result.avi'


def cut_frame(videoPath, saveFrameOccasionalPath):
    """
        将视频转化为图片帧，再经过检测处理
    """
    print("开始切帧！")
    video = cv2.VideoCapture(videoPath)
    fps = video.get(cv2.CAP_PROP_FPS)  # 获得原视频的fps
    frames = video.get(cv2.CAP_PROP_FRAME_COUNT)  # 获取原视频的总帧数

    print("fps={}, frames={}".format(int(fps), int(frames)))

    inters = tqdm(range(int(frames)))
    for i in inters:  # 遍历所有帧数
        ret, frame = video.read()

        # img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))  # 转化为PIL格式图片
        # frame = np.array(frame)
        # print(result_image.dtype)
        # result_image = cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)  # 转化回cv2格式
        cv2.imwrite(saveFrameOccasionalPath + str(i) + '.jpg', frame)  # 将每一帧进行保存
        inters.update(1)
    return fps


def detect(save_img=False):
    source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
    save_img = not opt.nosave and not source.endswith('.txt')  # save inference images
    webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
        ('rtsp://', 'rtmp://', 'http://', 'https://'))

    # Directories
    save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))  # increment run
    (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir

    # Initialize
    set_logging()
    device = select_device(opt.device)
    half = device.type != 'cpu'  # half precision only supported on CUDA

    # Load model
    model = attempt_load(weights, map_location=device)  # load FP32 model
    stride = int(model.stride.max())  # model stride
    imgsz = check_img_size(imgsz, s=stride)  # check img_size
    if half:
        model.half()  # to FP16

    # Second-stage classifier
    classify = False
    if classify:
        modelc = load_classifier(name='resnet101', n=2)  # initialize
        modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval()

    # Set Dataloader
    vid_path, vid_writer = None, None
    if webcam:
        view_img = check_imshow()
        cudnn.benchmark = True  # set True to speed up constant image size inference
        dataset = LoadStreams(source, img_size=imgsz, stride=stride)
        # print(dataset)
    else:
        dataset = LoadImages(source, img_size=imgsz, stride=stride)
        # print(dataset)

    # Get names and colors
    names = model.module.names if hasattr(model, 'module') else model.names
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]

    # Run inference
    if device.type != 'cpu':
        model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once
    t0 = time.time()
    for path, img, im0s, vid_cap in dataset:
        img = torch.from_numpy(img).to(device)
        img = img.half() if half else img.float()  # uint8 to fp16/32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        # Inference
        t1 = time_synchronized()
        pred = model(img, augment=opt.augment)[0]

        # Apply NMS
        pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
        t2 = time_synchronized()

        # Apply Classifier
        if classify:
            pred = apply_classifier(pred, modelc, img, im0s)

        # Process detections
        for i, det in enumerate(pred):  # detections per image
            if webcam:  # batch_size >= 1
                p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
            else:
                p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)

            p = Path(p)  # to Path
            save_path = str(save_dir / p.name)  # img.jpg
            txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # img.txt
            s += '%gx%g ' % img.shape[2:]  # print string
            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
            if len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                # Print results
                for c in det[:, -1].unique():
                    n = (det[:, -1] == c).sum()  # detections per class
                    class_ = names[int(c)]

                    s += f"{n} {class_}{'s' * (n > 1)}, "  # add to string

                # Write results
                for *xyxy, conf, cls in reversed(det):
                    if save_txt:  # Write to file
                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                        line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh)  # label format
                        with open(txt_path + '.txt', 'a') as f:
                            f.write(('%g ' * len(line)).rstrip() % line + '\n')

                    if save_img or view_img:  # Add bbox to image
                        name.append(names[int(cls)])  # 记录此类的名称
                        color.append(colors[int(cls)])  # 记录此类名称对应的颜色
                        label = f'{names[int(cls)]} {conf:.2f}'

                        plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
                # print(color)
                amount_name(name, color, im0, line_thickness=3)  # 在图上进行计数显示
                name.clear()  # 每判断完一张图片后，进行清空列表
                color.clear()  # 同上

            # Print time (inference + NMS)
            print(f'{s}Done. ({t2 - t1:.3f}s)')

            # Stream results
            if view_img:
                cv2.imshow(str(p), im0)
                cv2.waitKey(1)  # 1 millisecond

            # Save results (image with detections)
            if save_img:
                if dataset.mode == 'image':
                    cv2.imwrite(save_path, im0)
                else:  # 'frames' or 'stream'
                    if vid_path != save_path:  # new frames
                        vid_path = save_path
                        if isinstance(vid_writer, cv2.VideoWriter):
                            vid_writer.release()  # release previous frames writer
                        if vid_cap:  # frames
                            fps = vid_cap.get(cv2.CAP_PROP_FPS)
                            w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                            h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                        else:  # stream
                            fps, w, h = 30, im0.shape[1], im0.shape[0]
                            save_path += '.mp4'
                        vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
                    vid_writer.write(im0)

    if save_txt or save_img:
        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
        print(f"Results saved to {save_dir}{s}")
    # print(save_dir)  # 结果保存的路径
    print(f'Done. ({time.time() - t0:.3f}s)')
    save_result = save_dir
    return save_result


def integration_frame(saveFrameOccasionalPath, save_path, fps):
    """
        将图片帧转化为视频
    """
    im_list = os.listdir(saveFrameOccasionalPath)  # 获取帧数文件夹里面图片文件的名字
    im_list.sort(key=lambda x: int(x.replace("frame", "").split('.')[0]))  # 查看图片顺序对不对

    img = Image.open(os.path.join(saveFrameOccasionalPath, im_list[0]))  # 取第一帧
    img_size = img.size  # 得到图片大小

    fourcc = cv2.VideoWriter_fourcc(*'XVID')  # 进行编码
    videoWrite = cv2.VideoWriter(save_path, fourcc, fps, img_size)  # 写入视频的对象

    for i in tqdm(im_list):  # 遍历帧数名列表
        im_name = os.path.join(saveFrameOccasionalPath + '/' + i)
        frame = cv2.imdecode(np.fromfile(im_name, dtype=np.uint8), -1)  # 得到编码后的帧
        videoWrite.write(frame)  # 整合为视频

    videoWrite.release()  # 释放原视频


def remove_file(path):
    print("正在删除冗余文件.......")
    imgList = os.listdir(path)

    num_imgs = len(imgList)
    # count = 0
    inters = tqdm(range(num_imgs))
    for i in inters:
        imgPath = os.path.join(path, imgList[i])

        os.remove(imgPath)
        inters.update(1)
    print("删除完成！")


if __name__ == '__main__':
    fps = cut_frame(video_path, frame_occasional_save)

    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default='runs/train/exp22/weights/best.pt',
                        help='model.pt path(s)')
    parser.add_argument('--source', type=str, default='data/images/frames', help='source')  # file/folder, 0 for webcam
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', help='display results')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--update', action='store_true', help='update all models')
    parser.add_argument('--project', default='runs/detect', help='save results to project/name')
    parser.add_argument('--name', default='exp', help='save results to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    opt = parser.parse_args()
    print(opt)
    check_requirements(exclude=('pycocotools', 'thop'))

    with torch.no_grad():
        if opt.update:  # update all models (to fix SourceChangeWarning)
            for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
                save_ = detect()
                strip_optimizer(opt.weights)
        else:
            save_ = detect()
    a = str(save_).replace("\\", "/")
    integration_frame(first_path + a, video_result_path, fps)
    remove_file(frame_occasional_save)
    remove_file(save_)
